import pandas as pd
from datetime import datetime
import plotly.graph_objects as go

df = pd.read_csv('新冠肺炎疫情历史数据-部分.csv')
# df.info()
df['时间'] = df['时间'].apply(lambda d: datetime.strptime(d, '%Y.%m.%d'))
df = df.sort_values(by=['时间'])
sankey_data = df.groupby(['省份/直辖市']).last().reset_index()
print(sankey_data)

# 市场列表
node = []
# 列表生成式
pro = sankey_data['省份/直辖市'].unique().tolist()
node += [{'label': prov} for prov in pro ]
node += [{'label': sankey_data['时间'][0]}, {'label': '治愈'}, {'label': '死亡'}]
print(node)
#链接列表
links = []
for index, p in enumerate(pro):

    rgba = f'rgba({random})'
    links.append({
        "source":index,
        "target":len(pro),
        "value" :sankey_data[sankey_data['省份/直辖市']==p]['总确诊'].values[0]

    })
    #时间 ---（治愈人数）--治愈
    links.append({
        "source": len(pro),
        "target": len(pro)+1,
        "value": sankey_data[sankey_data['省份/直辖市'] == p]['治愈'].values[0],
        "color":rgba
    })
    #时间 ---（死亡人数）--死亡
    links.append({
        "source": index,
        "target": len(pro)+2,
        "value": sankey_data[sankey_data['省份/直辖市'] == p]['死亡'].values[0],
        "color":rgba
    })

print(len(links))

#图纸
fig = go.Figure()
fig.update_layout(title_text='总确诊的治愈和死亡情况')
fig.add_sankey(
    node=dict(label=[n["label"]for n in node]),
    link=dict(
        source=[link['source'] for link in links],
        target=[link['target'] for link in links],
        value=[link['value'] for link in links],
        color=[link['color'] for link in links]
    )
)
fig.show()